import shutil, os, nibabel

import PIL.Image


def niito2D(filepath , suffix = ".png" ,resize_shape=(256,256), sample_count = 3 , start_ratio=0 , end_ratio=1):
    # left = int(count_per_nii/2)
    # right = int(count_per_nii/2)+1
    inputfiles = os.listdir(filepath)  # 遍历文件夹数据
    # outputfile = 'slices/axial150x256/'  # 输出文件夹
    outputpath = filepath + '/output'
    # set destination folder
    if not os.path.exists(outputpath):
        os.makedirs(outputpath)  # 不存在输出文件夹则新建
        print("Created ouput directory: " + outputpath)
    print('Input file is ', inputfiles)
    print('Output folder is ', outputpath)
    i=0
    for inputfile in inputfiles:
        files = os.listdir(filepath+"/" + inputfile)
        files = list(filter(lambda f: f.endswith('.nii'), files))
        for f in files:
            fp = filepath+"/" + inputfile + "/" + f
            try:
                nii = nibabel.load(fp)  # 数据读取
                image_array = nii.get_fdata()
            except Exception as e:
                print(e)
                continue
            shape = image_array.shape
            print(shape)

            #中间集中采样
            # l = shape[-1]
            # if int(l/6)-left >=0 :
            #     min = (int(l/6)-left)
            # else:
            #     min =0
            # if int(l*5/6)+right <=l :
            #     max =int(l*5/6)+right
            # else:
            #     max =l

            # intervals = ((min,int(l/6)+left) , (int(l/2) - left , int(l/2)+right) , (int(l*5/6)-left,max))
            # for interval in intervals:
            length = shape[2]
            center = int(length / 2)
            if sample_count %2 ==0:
                count_per_side = sample_count / 2
                left_sample = int((0.5 - start_ratio)*length)
                right_sample = int((end_ratio- 0.5)*length)
                left_step = -int(left_sample / (count_per_side-1))
                right_step = int(right_sample / (count_per_side+1))
                intervals =  ((center+right_step , int(length * end_ratio) , right_step), (center +left_step , int(length * start_ratio), left_step))
            else:
                count_per_side = (sample_count-1) / 2
                left_sample = int((0.5 - start_ratio)*length)
                right_sample = int((end_ratio- 0.5)*length)
                left_step = -int(left_sample / (count_per_side+1))-1
                right_step = int(right_sample / (count_per_side+1))+1
                intervals = ((center,center+1,1),(center+right_step , int(length * end_ratio) , right_step), (center +left_step , int(length * start_ratio), left_step))

            for interval in intervals:
                for slice in range(interval[0] , interval[1] , interval[2]):
                        cur_img = image_array[:, : ,slice]
                        image_name = str(i)+suffix
                        #数据增强,并直接保存至mat
                        cur_img = PIL.Image.fromarray(cur_img/(cur_img.max() - cur_img.min())*255).convert('L')
                        if not (resize_shape[0] == cur_img.size[0] and resize_shape[1] == cur_img.size[1]):
                            cur_img = cur_img.resize(resize_shape, PIL.Image.Resampling.BICUBIC)
                        cur_img.save(os.path.join(outputpath , image_name))
                        i+=1


    print('Finished converting images')




if __name__ == '__main__':
    # 获取该路径下所有三维图像压缩包的切片图像
    # niito2D("E:\download\dataset\keras\IXI-T1" , start_ratio=0.2,end_ratio=0.8 , sample_count=7)
    # niito2D("E:\download\dataset\keras\IXI-T2" , end_ratio=0.8, sample_count=7)
    # niito2D("E:\download\dataset\keras\IXI-PD" ,start_ratio=0.2, sample_count=7)
    niito2D("E:\download\dataset\keras\IXI-MRA", sample_count=8)
    niito2D("E:\download\dataset\keras\IXI-DTI",sample_count=1)

